#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
特征工程脚本 - 使用origin原始数据
Feature Engineering Script - Using Origin Raw Data
"""

import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder
import os

def handle_missing_values(df, is_train=True, train_df=None):
    """处理缺失值 - 分别处理训练和测试集"""
    if is_train:
        print("  训练集缺失值处理...")
        initial_shape = df.shape
        print(f"    初始形状: {initial_shape}")

        # 1. 删除缺失值过多的列（缺失值超过50%）
        missing_ratio = df.isnull().mean()
        high_missing_cols = missing_ratio[missing_ratio > 0.5].index.tolist()

        # 2. 删除指定的缺失值过多的列
        cols_to_drop = ['Alley', 'Fence', 'FireplaceQu', 'PoolQC', 'MasVnrType', 'MiscFeature']
        all_drop_cols = list(set(high_missing_cols + cols_to_drop))
        existing_cols = [col for col in all_drop_cols if col in df.columns]

        if existing_cols:
            print(f"    删除高缺失值列: {len(existing_cols)}个")
            df = df.drop(columns=existing_cols)

        # 3. 删除缺失值较多的行
        columns_to_check = ['LotFrontage', 'BsmtQual', 'BsmtCond', 'BsmtExposure',
                            'BsmtFinType1', 'BsmtFinType2', 'GarageType', 'GarageYrBlt',
                            'GarageFinish', 'GarageQual', 'GarageCond', 'MasVnrArea']

        existing_check_cols = [col for col in columns_to_check if col in df.columns]

        if existing_check_cols:
            rows_before = len(df)
            df = df.dropna(subset=existing_check_cols)
            rows_after = len(df)
            print(f"    删除缺失值行: {rows_before - rows_after} 行")

        # 4. 数值列用中位数填充
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        missing_numeric_cols = [col for col in numeric_cols if df[col].isnull().any()]

        if missing_numeric_cols:
            for col in missing_numeric_cols:
                median_val = df[col].median()
                df[col] = df[col].fillna(median_val)

        # 5. 分类型列用众数填充
        categorical_cols = df.select_dtypes(include=['object', 'category']).columns
        missing_cat_cols = [col for col in categorical_cols if df[col].isnull().any()]

        if missing_cat_cols:
            for col in missing_cat_cols:
                mode_val = df[col].mode()[0] if not df[col].mode().empty else 'Unknown'
                df[col] = df[col].fillna(mode_val)

        final_shape = df.shape
        print(f"    最终形状: {final_shape}")
        print(f"    减少行数: {initial_shape[0] - final_shape[0]}")
        print(f"    减少列数: {initial_shape[1] - final_shape[1]}")

    else:
        print("  测试集缺失值处理...")
        initial_shape = df.shape
        print(f"    初始形状: {initial_shape}")

        # 删除与训练集相同的列
        cols_to_drop = ['Alley', 'Fence', 'FireplaceQu', 'PoolQC', 'MasVnrType', 'MiscFeature']
        existing_cols = [col for col in cols_to_drop if col in df.columns]

        if existing_cols:
            print(f"    删除指定列: {len(existing_cols)}个")
            df = df.drop(columns=existing_cols)

        # 使用训练集的统计量填充数值列
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        for col in numeric_cols:
            if df[col].isnull().any() and col in train_df.columns:
                median_val = train_df[col].median()
                df[col] = df[col].fillna(median_val)

        # 使用训练集的众数填充分类列
        categorical_cols = df.select_dtypes(include=['object', 'category']).columns
        for col in categorical_cols:
            if df[col].isnull().any() and col in train_df.columns:
                mode_val = train_df[col].mode()[0] if not train_df[col].mode().empty else 'Unknown'
                df[col] = df[col].fillna(mode_val)

        final_shape = df.shape
        print(f"    最终形状: {final_shape}")

    return df

def addFeatures(df):
    """构造新特征"""
    print("构造新特征...")

    # 基础特征组合
    df['YrBltAndRemod'] = df['YearBuilt'] + df['YearRemodAdd']
    df['TotalSF'] = df['TotalBsmtSF'] + df['1stFlrSF'] + df['2ndFlrSF']
    df['Total_sqr_footage'] = df['BsmtFinSF1'] + df['BsmtFinSF2'] + df['1stFlrSF'] + df['2ndFlrSF']
    df['Total_Bathrooms'] = df['FullBath'] + (0.5 * df['HalfBath']) + df['BsmtFullBath'] + (0.5 * df['BsmtHalfBath'])
    df['Total_porch_sf'] = df['OpenPorchSF'] + df['3SsnPorch'] + df['EnclosedPorch'] + df['ScreenPorch'] + df['WoodDeckSF']

    # 额外特征
    df['HouseAge'] = df['YrSold'] - df['YearBuilt']
    df['RemodAge'] = df['YrSold'] - df['YearRemodAdd']
    df['TotalSF_per_room'] = df['TotalSF'] / df['TotRmsAbvGrd']
    df['LivingArea_ratio'] = df['GrLivArea'] / df['TotalSF']

    return df

def check_final_missing(df, stage_name="处理后"):
    """检查最终缺失值情况"""
    print(f"\n{stage_name}缺失值检查:")

    total_missing = df.isnull().sum().sum()
    missing_cols = df.isnull().any()
    cols_with_missing = missing_cols[missing_cols].index.tolist()

    print(f"  总缺失值数量: {total_missing}")
    print(f"  有缺失值的列数: {len(cols_with_missing)}")

    if cols_with_missing:
        print("  缺失值最多的5列:")
        missing_counts = df.isnull().sum().sort_values(ascending=False)
        top_missing = missing_counts.head(5)
        for col, count in top_missing.items():
            print(f"    {col}: {count}")
    else:
        print("  [OK] 没有缺失值")

    return total_missing, len(cols_with_missing)

def encoding(df):
    """分类变量独热编码"""
    print("进行独热编码...")

    categorical_cols = list(df.select_dtypes(include='object').columns)
    numerical_cols = list(df.select_dtypes(include='number').columns)

    print(f"分类变量: {len(categorical_cols)}")
    print(f"数值变量: {len(numerical_cols)}")

    if not categorical_cols:
        return df

    ohe = OneHotEncoder(sparse_output=False, drop='first', handle_unknown='ignore')
    encoded_values = ohe.fit_transform(df[categorical_cols])
    encoded_cols = list(ohe.get_feature_names_out())

    df_encoded = pd.DataFrame(encoded_values, columns=encoded_cols, index=df.index)
    df_numerical = df[numerical_cols]
    df_new = pd.concat([df_numerical, df_encoded], axis=1)

    return df_new

def main():
    """主处理流程"""
    print("="*50)
    print("特征工程V2 - 分开处理训练和测试集")
    print("="*50)

    # 1. 加载原始数据
    print("1. 加载原始数据...")
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_dir = os.path.dirname(script_dir)

    train_path = os.path.join(project_dir, 'origin', 'train.csv')
    test_path = os.path.join(project_dir, 'origin', 'test.csv')

    train_df = pd.read_csv(train_path)
    test_df = pd.read_csv(test_path)

    print(f"训练数据: {train_df.shape}")
    print(f"测试数据: {test_df.shape}")

    # 2. 保存原始目标变量
    target = train_df['SalePrice'].copy()
    print(f"目标变量形状: {target.shape}")

    # 3. 分别处理训练和测试集的缺失值
    print("2. 处理训练集缺失值...")
    train_df = handle_missing_values(train_df, is_train=True)

    print("3. 处理测试集缺失值...")
    test_df = handle_missing_values(test_df, is_train=False, train_df=train_df)

    # 4. 构造新特征
    print("4. 构造新特征...")
    train_df = addFeatures(train_df)
    test_df = addFeatures(test_df)
    print(f"特征构造后 - 训练集: {train_df.shape}, 测试集: {test_df.shape}")

    # 5. 独热编码（基于训练集拟合）
    print("5. 独热编码...")
    categorical_cols = list(train_df.select_dtypes(include='object').columns)
    numerical_cols = list(train_df.select_dtypes(include='number').columns)

    # 从训练集和测试集的数值列中移除Id和SalePrice
    numerical_cols = [col for col in numerical_cols if col not in ['Id', 'SalePrice']]

    if categorical_cols:
        ohe = OneHotEncoder(sparse_output=False, drop='first', handle_unknown='ignore')

        # 在训练集上拟合
        train_cat = ohe.fit_transform(train_df[categorical_cols])
        test_cat = ohe.transform(test_df[categorical_cols])

        encoded_cols = list(ohe.get_feature_names_out())

        # 创建编码后的DataFrame
        train_encoded = pd.DataFrame(train_cat, columns=encoded_cols, index=train_df.index)
        test_encoded = pd.DataFrame(test_cat, columns=encoded_cols, index=test_df.index)

        # 合并数值和编码特征
        train_final = pd.concat([train_df[numerical_cols], train_encoded], axis=1)
        test_final = pd.concat([test_df[numerical_cols], test_encoded], axis=1)
    else:
        # 没有分类变量
        train_final = train_df[numerical_cols].copy()
        test_final = test_df[numerical_cols].copy()

    # 6. 最终缺失值清理
    print("6. 最终缺失值清理...")
    train_final = train_final.fillna(0)
    test_final = test_final.fillna(0)

    # 7. 保存完整数据集
    print("7. 保存文件...")
    output_dir = os.path.join(project_dir, 'processed_data', 'v2')
    os.makedirs(output_dir, exist_ok=True)

    # 保存完整数据集 - 修复索引对齐问题
    # 确保target与处理后的train_df索引对齐
    target_filtered = target.iloc[train_df.index]
    train_full = pd.DataFrame({'Id': train_df['Id'], 'SalePrice': target_filtered}, index=train_df.index)
    train_full = pd.concat([train_full, train_final], axis=1)
    train_full.to_csv(os.path.join(output_dir, 'train_full_v2.csv'), index=False)

    test_full = pd.DataFrame({'Id': test_df['Id']}, index=test_df.index)
    test_full = pd.concat([test_full, test_final], axis=1)
    test_full.to_csv(os.path.join(output_dir, 'test_full_v2.csv'), index=False)

    print(f"\n最终结果:")
    print(f"训练集: {train_full.shape}")
    print(f"测试集: {test_full.shape}")
    print(f"特征数量: {train_final.shape[1]}")
    print(f"无缺失值: 训练集{train_final.isnull().sum().sum()}, 测试集{test_final.isnull().sum().sum()}")

    print(f"\n文件已保存到 {output_dir}:")
    print("- train_full_v2.csv: 完整训练集")
    print("- test_full_v2.csv: 完整测试集")

    return {
        'train_full': train_full,
        'test_full': test_full,
        'X_train': train_final,
        'X_test': test_final,
        'y_train': target
    }

if __name__ == "__main__":
    main()